import itertools
from itertools import islice
from random import randrange

from number_theory.gcd import inv_mod
from polynomial.gcd import gcd
from polynomial.polynomial_Zn import PolynomialZn

__all__ = ["factor", "factor_unique"]



def factor(poly):
    if not poly:
        return [(type(poly)(), 1)]
    leading_coeff = poly.leading_coeff
    if leading_coeff == 1:
        result = []
    else:
        result = [(type(poly)(leading_coeff), 1)]
    for p, deg in square_free(poly):
        result += ((f, deg) for f in factor_unique(p))
    return sorted(result)


def square_free(poly):
    PolyType = type(poly)
    prime = PolyType.mod
    one = PolyType(1)
    result = [(one, 0)]

    f = poly.monic()
    h = PolyType()
    for deg in itertools.count(0):
        d = f.deriv()
        f1 = gcd(f, d)
        h1 = f // f1
        g = h // h1
        if deg and g.deg > 0:
            result.append((g, deg))
        f = f1
        h = h1
        if h.deg == 0:
            break
    
    if f.deg != 0:
        spoly = f[::prime]
        sfactor = square_free(spoly)
        k = len(result)
        l = len(sfactor)
        for j, (p, d) in enumerate(sfactor):
            result.append((p, d * prime))
            head = (j + 1) * k
            for i in range(1, k):
                g = gcd(result[i][0], result[head][0])
                result.append((g, result[i][1] + result[head][1]))
                result[i] = result[i][0] // g, result[i][1]
                result[head] = result[head][0] // g, result[head][1]

    return list(filter(lambda x: x[0] != one, result))


def factor_unique(poly):
    result = []
    for p, deg in distinct_degree(poly):
        result += equal_degree(p, deg)
    return result


def berlekamp(poly):
    if poly.deg == 1:
        return [poly]

    PolyType = type(poly)
    result = []
    prime = PolyType.mod
    degree = poly.deg
    q = [[0] * degree for x in range(degree)]
    base_poly = pow(PolyType((0, 1)), prime, poly)
    p = PolyType(1)
    for i in range(degree):
        for j in range(degree):
            q[i][j] = p[j]
        p = p * base_poly % poly
    for i in range(degree):
        q[i][i] = (q[i][i] - 1) % prime
    base_set = set(range(degree))
    for c in range(degree):
        m = 1
        while m < degree and q[m][c] == 0:
            m += 1
        if m < degree:
            base_set.discard(m)
            inverse = inv_mod(q[m][c], prime)
            for j in filter(lambda x: x != c and q[m][x], range(degree)):
                temp = q[m][j] * inverse % prime
                for i in range(1, degree):
                    q[i][j] = (q[i][j] - temp * q[i][c]) % prime
    basis = sorted(list(base_set))
    rank = len(basis)

    base = []
    count = 0
    for i in range(degree):
        if i in base_set:
            b = [0] * rank
            b[count] = 1
            count += 1
        else:
            column = 0
            while q[i][column] == 0:
                column += 1
            inverse = prime - inv_mod(q[i][column], prime)
            b = [q[x][column] * inverse % prime for x in basis]
        base.append(b)

    def trial():
        nonlocal length
        p = PolyType(base[j][i] for j in range(degree))
        for j in range(prime):
            p[0] = j
            for x in range(length):
                g = gcd(p, f[x])
                if 0 < g.deg < f[x].deg:
                    f.append(f[x] // g)
                    f[x] = g
                    length += 1
                    if length == rank:
                        return

    f = [poly]
    length = 1
    for i in range(1, rank):
        trial()
    result += f
    return result


def fast_pow_mod(poly, prime, mod):
    for deg in range(poly.deg, 0, -1):
        poly[deg * prime] = poly[deg]
        poly[deg] = 0
    poly %= mod


def distinct_degree(poly):
    PolyType = type(poly)
    prime = PolyType.mod
    f = PolyType(poly)
    base = PolyType((0, 1))
    h = PolyType(base)
    result = []
    for degree in itertools.count(1):
        if f.deg == 0:
            break
        elif degree << 1 > f.deg:
            result.append((f, f.deg))
            break
        else:
            if prime <= 3:
                fast_pow_mod(h, prime, poly)
            else:
                h = pow(h, prime, poly)
            g = gcd(h - base, f)
            if g.deg > 0:
                result.append((g, degree))
                f //= g
    return result


def equal_degree(poly, degree):
    PolyType = type(poly)
    prime = PolyType.mod
    result = []

    if prime != 2:
        exp = (prime ** degree) >> 1
    stack = [poly]
    while stack:
        top = stack.pop()
        top_degree = top.deg
        if top_degree == degree:
            result.append(top)
            continue
        
        while True:
            a = PolyType()
            while a.deg <= 0:
                a = PolyType(randrange(prime) for x in range(top_degree))
            g = gcd(a, top)
            if g.deg > 0:
                stack += g, top // g
                break
            if prime == 2:
                b = PolyType(a)
                for d in range(degree - 1):
                    fast_pow_mod(a, 2, top)
                    b += a
            else:
                b = pow(a, exp, top)
            b[0] -= 1
            g = gcd(b, top)
            if 0 < g.deg < top_degree:
                stack += g, top // g
                break

    return result
